# functions used to prepare the CUDA kernel

from misc import read_file, write_file, find_index
from numpy import *
from numpy.random import rand
from pycuda.compiler import SourceModule
import sys
from texture import *
from param import Param
from dataset import Dataset


typenames = {'FLOAT': 'float', 'FLOAT_RANGE': 'float',
    'FLOAT2': 'float2', 'FLOAT2_RANGE': 'float2',
    'FLOAT3': 'float3', 'FLOAT3_RANGE': 'float3',
    'FLOAT4': 'float4', 'FLOAT4_RANGE': 'float4'}
    
    
class Kernel:
    # members: shader, full_shader, params, datasets, module, run

    def __init__(self, shader_file, options):
        self.param_resolution = int(options.param_resolution)
    
        self.shader = read_file(shader_file)
        self.parse_params()

        # add size parameters
        for d in self.datasets:
            self.params.append(Param(d.name + '_size', 'float' + str(d.dimension), None))
            
        # generate full shader    
        tex_defs = self.generate_texture_defs()
        members = self.generate_members()
        init = self.generate_initializer()
        self.full_shader = read_file(sys.path[0] + "/common.cu")
        self.full_shader += read_file(sys.path[0] + "/vr.cu") % (tex_defs,
            members, members, init, self.shader)
        write_file("tmp", self.full_shader)
        

    def load(self):
        # create CUDA source module
        self.module = SourceModule(self.full_shader, no_extern_c = True,
            options = ["-use_fast_math", "-O3"])
        self.run = self.module.get_function("run")

        # load data and fix size parameters
        for d in self.datasets:
            d.load(self.module)
            idx = find_index(lambda x: x.name == d.name + '_size', self.params)
            self.params[idx].value = d.size
            
        self.create_helper_textures()
        
    
    def parse_params(self):
        "Parse parameter definitions from shader source"

        def conv(p):
            xs = p.split()
            if len(xs) == 1:
                return float(xs[0])
            else:
                return tuple(map(float, xs))

        self.params = []
        self.datasets = []

        for line in self.shader.splitlines():
            line = line.strip()
            
            if line.startswith('DATA') or line.startswith('PARAM'):
                paren1 = line.find("(")
                paren2 = line.rfind(")")
                cmd = line[:paren1]
                args = line[paren1+1:paren2].split(",")
                args = [x.strip("'\" ") for x in args]
                
                if cmd == 'DATA':
                    # Dataset args: name, filename, typename, dimension, num_channels
                    self.datasets.append(Dataset(args[0], args[1], args[2], int(args[3]), int(args[4])))
                elif cmd == 'PARAM':
                    # Param args: name, typename, value
                    self.params.append(Param(args[0], 'float', float(args[1])))
                elif cmd == 'PARAM_RANGE':
                    # Param args: name, typename, value, bounds=None, key=None
                    self.params.append(Param(args[0], 'float', float(args[1]),
                        (float(args[2]), float(args[3])), args[4], self.param_resolution))
                else:
                    assert False
            

    def generate_members(self):
        code = ""
        for p in self.params:
            code += "%s %s;\n" % (p.typename, p.name)
        return code


    def generate_initializer(self):
        code = "__device__ void init(Params* params)\n{\n"
        for p in self.params:
            code += "    %s = params->%s;\n" % (p.name, p.name)
        code += "}\n"
        return code


    def generate_getters(self):
        code = ""
        for p in self.params:
            code += "__device__ %s %s() { return params->%s; }\n" % (p.typename, p.name, p.name)
        return code


    def generate_texture_defs(self):
        code = ""
        
        for d in self.datasets:
            t = d.typename
            if d.num_channels == 4: t += '4'

            e = 'cudaReadModeNormalizedFloat'
            if d.typename == 'float': e = 'cudaReadModeElementType'
            
            if d.dimension <= 3:
                code += "texture<%s, %s, %s> %s;\n" % (t, d.dimension, e, d.name)
            else:
                # there are no 4D textures in CUDA, so have to use 3D
                code += "texture<%s, 3, %s> %s;\n" % (d.typename, e, d.name)

        return code


    def create_helper_textures(self):
        "create lookup textures for cubic interpolation and random generation"
        
        def hg(a):
            a2 = a * a
            a3 = a2 * a
            w0 = (-a3 + 3*a2 - 3*a + 1) / 6
            w1 = (3*a3 - 6*a2 + 4) / 6
            w2 = (-3*a3 + 3*a2 + 3*a + 1) / 6
            w3 = a3 / 6
            g = w2 + w3
            h0 = 1 - w1 / (w0 + w1) + a
            h1 = 1 + w3 / (w2 + w3) - a
            return h0, h1, g, 0
            
        def dhg(a):
            a2 = a * a
            w0 = (-a2 + 2*a - 1) / 2
            w1 = (3*a2 - 4*a) / 2
            w2 = (-3*a2 + 2*a + 1) / 2
            w3 = a2 / 2
            g = w2 + w3
            h0 = 1 - w1 / (w0 + w1) + a
            h1 = 1 + w3 / (w2 + w3) - a
            return h0, h1, g, 0

        tmp = zeros((256, 4), dtype=float32)
        for i, x in enumerate(linspace(0, 1, 256)):
            tmp[i, :] = array(hg(x))
        tmp = reshape(tmp, (1, 256, 4))
        hg_texture = create_2d_rgba_texture(tmp, self.module, 'hgTexture')

        tmp = zeros((256, 4), dtype=float32)
        for i, x in enumerate(linspace(0, 1, 256)):
            tmp[i, :] = array(dhg(x))
        tmp = reshape(tmp, (1, 256, 4))
        dhg_texture = create_2d_rgba_texture(tmp, self.module, 'dhgTexture')
        
        tmp = uint32(rand(256, 256) * (2 << 30))
        random_texture = create_2d_texture(tmp, self.module, 'randomTexture', True)
        
        # to prevent GC from destroying the textures
        self.dummy = (hg_texture, dhg_texture, random_texture)
    
    
    def create_params_struct(self):
        """
    	Create the Params structure to be passed to the CUDA kernel.
    	Make sure alignment is handled correctly.
    	"""

        align = [0, 1, 2, 1, 4]
        l = []
        
        for p in self.params:
            # print p.name, p.value
            if isinstance(p.value, tuple):
                v = list(p.value)
            else:
                v = [p.value]

            # align if needed
            while len(l) % align[len(v)] != 0:
    	        l.append(float(0))

            l.extend(v)

        l.append(1234.0) # dummy, to ensure the struct is non-empty
        # print 'pstruct:', l
        return array(l, dtype = float32)
